# -*- coding: utf-8 -*-

# __date:       2021/6/11
# __author:     YangChao
# __function:   绘制图片

import os
import matplotlib.pyplot as plt


def pic_plot(data,
             name: str,
             save_config: str,
             best_index: int = None,
             xlabel: str = 'Points',
             ylabel: str = 'Leak',
             title: str = 'calibration_porcess',
             ):
    """绘制折线图形"""
    font = {'family': 'Times New Roman',
            'weight': 'bold',
            'size': 20,
            }
    x = range(0, len(data))
    fig, ax = plt.subplots(figsize=(14, 12))
    ax.set_xlabel(xlabel, fontdict=font)
    ax.set_ylabel(ylabel, fontdict=font)
    ax.set_title(name, fontdict=font)
    ax.plot(data, label='leak')
    ax.set_xticks(x)
    ax.legend(loc=4, fontsize=20)

    _add_annotate(ax, best_index, data[best_index])

    pic_name = f'/{name}.png'
    path = save_config + '/picture'
    if not os.path.exists(path):
        os.makedirs(path)
    fig.savefig(path + pic_name)
    plt.close(fig)


def pic_scatter(x,
                y,
                z,
                name: str,
                param: str,
                save_config: str,
                best_index: int = None,
                ):
    """绘制散点图"""
    fontc = {'family': 'serif',
             'color': 'darkred',
             'weight': 'normal',
             'size': 20}
    font = {'family': 'Times New Roman',
            'weight': 'bold',
            'size': 20}
    labels = param.split('-')
    fig, ax = plt.subplots(figsize=(14, 12))
    ax.set_xlabel(labels[0].upper(), fontdict=font)
    ax.set_ylabel(labels[1].upper(), fontdict=font)
    ax.set_title(name, fontdict=font)
    my_ax = ax.scatter(x=x, y=y, c=z, cmap="Blues")
    cb = fig.colorbar(my_ax, ax=ax)
    cb.set_label('Leak value', fontdict=fontc)
    _add_annotate(ax, x[best_index], y[best_index])

    pic_name = f'/{name}.png'
    path = save_config + '/picture'
    if not os.path.exists(path):
        os.makedirs(path)
    fig.savefig(path + pic_name)
    plt.close(fig)


def _add_annotate(ax, a, b):
    """添加注解"""
    bbox = dict(boxstyle="round, pad=1.5", fc="yellow")
    arrowprops = dict(arrowstyle="->", connectionstyle="arc3", color='lawngreen')
    ax.scatter(a, b, marker='*', c='lawngreen', linewidths=10)
    ax.annotate('Good point({:.2f}, {:.2f})'.format(a, b),
                xy=(a, b),
                textcoords='figure fraction',
                xytext=(0.2, 0.8),
                weight='extra bold',
                bbox=bbox,
                arrowprops=arrowprops
                )
